from mxnet import gluon
import mxnet as mx
import d2lzh as d2l

batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size=batch_size)

ctx=mx.gpu()

net = gluon.nn.Sequential()
with net.name_scope():
    net.add(gluon.nn.Conv2D(channels = 20, kernel_size = 5, activation = 'relu'))
    net.add(gluon.nn.MaxPool2D(pool_size = 2, strides = 2))
    net.add(gluon.nn.Conv2D(channels = 50, kernel_size = 3, activation = 'relu'))
    net.add(gluon.nn.MaxPool2D(pool_size = 2, strides = 2))
    net.add(gluon.nn.Flatten())
    net.add(gluon.nn.Dense(128, activation = 'relu'))
    net.add(gluon.nn.Dense(10))
net.initialize(force_reinit=True, ctx=ctx)

trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.1})

d2l.train_ch5(net, train_iter, test_iter, 10, ctx, 256, None, None, trainer)
